import io
import sys
import traceback
import pandas as pd
import numpy as np
import os
import random as rn
import tensorflow as tf
import logging
from model.DGCN import get_DGCN_Model
from model.utils import convert_combinations_to_ids, match_ids_with_data, get_adj_mats, extract_feature_in_x_all

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

PRO_BASE_DIR = os.path.dirname(__file__)
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

def start(file_path):
    output = io.StringIO()
    sys.stdout = output
    tableData = []
    log_messages = []
    try:
        test_df = pd.read_csv(file_path)
        disease_id_df = pd.read_csv(os.path.join(PRO_BASE_DIR, 'disease_id.csv'))
        drug_id_df = pd.read_csv(os.path.join(PRO_BASE_DIR, 'drug_id.csv'))

        valid_combinations = []
        for _, row in test_df.iterrows():
            if row['disease_name'] not in disease_id_df['dis_name'].values or row['drug_name'] not in drug_id_df['drug_name'].values:
                msg = f"Sorry, the prediction task for the association between the parasitic disease {row['disease_name']} and the drug {row['drug_name']} is not supported yet."
                logging.info(msg)
                log_messages.append(msg)
            else:
                valid_combinations.append(row)

        if not valid_combinations:
            return output.getvalue(), tableData, log_messages

        valid_combinations_df = pd.DataFrame(valid_combinations)
        converted_combinations_df = convert_combinations_to_ids(valid_combinations_df, disease_id_df, drug_id_df)

        if converted_combinations_df.empty:
            for _, row in valid_combinations_df.iterrows():
                msg = f"Sorry, the prediction task for the association between the parasitic disease {row['disease_name']} and the drug {row['drug_name']} is not supported yet."
                logging.info(msg)
                log_messages.append(msg)
            return output.getvalue(), tableData, log_messages

        mode = 0
        fold = 0
        train_file_path = os.path.join(PRO_BASE_DIR, f"mode{mode}_fold{fold}_X_train.csv")
        test_file_path = os.path.join(PRO_BASE_DIR, f"mode{mode}_fold{fold}_X_test.csv")

        matched_data = match_ids_with_data(train_file_path, test_file_path, converted_combinations_df)

        if not matched_data:
            for _, row in converted_combinations_df.iterrows():
                msg = f"Sorry, the prediction task for the association between the parasitic disease {row['disease_name']} and the drug {row['drug_name']} is not supported yet."
                logging.info(msg)
                log_messages.append(msg)
            return output.getvalue(), tableData, log_messages

        SEED = 123
        os.environ['PYTHONHASHSEED'] = str(SEED)
        os.environ['TF_DETERMINISTIC_OPS'] = '0'
        tf.random.set_seed(SEED)
        np.random.seed(SEED)
        rn.seed(SEED)

        epoch = 1712
        lr = 0.1
        bs = 256
        emb_dim = 128

        NUM_ENTITIES = 197
        NUM_RELATIONS = 3

        OUTPUT_DIM = emb_dim
        X_train = pd.read_csv(train_file_path)
        X_test = pd.read_csv(test_file_path)
        x_all = pd.read_csv(os.path.join(PRO_BASE_DIR, f"x_all_{emb_dim}.csv"), header=0).values

        model = get_DGCN_Model(
            num_entities=NUM_ENTITIES,
            num_relations=NUM_RELATIONS,
            embedding_dim=emb_dim,
            output_dim=OUTPUT_DIM,
            seed=SEED,
            all_feature_matrix=x_all,
            mode=mode,
            fold=fold
        )

        model.load_weights(os.path.join(PRO_BASE_DIR, f'mode{mode}_fold{fold}_epoch{epoch}_learnRate{lr}_batchsize{bs}_embdim{emb_dim}_layer2.h5'))

        for _, row in converted_combinations_df.iterrows():
            matched_data = match_ids_with_data(train_file_path, test_file_path, pd.DataFrame([row]))

            if not matched_data:
                msg = f"Sorry, the prediction task for the association between the parasitic disease {row['disease_name']} and the drug {row['drug_name']} is not supported yet."
                logging.info(msg)
                log_messages.append(msg)
                continue

            match_train, match_test, drug_name, disease_name = matched_data[0]

            if match_train.empty and match_test.empty:
                msg = f"Sorry, the prediction task for the association between the parasitic disease {disease_name} and the drug {drug_name} is not supported yet."
                logging.info(msg)
                log_messages.append(msg)
                continue

            match = pd.concat([match_train, match_test])
            x_all_test = extract_feature_in_x_all(match, x_all)

            ADJACENCY_DATA = tf.concat([X_train, X_test, match], axis=0)
            ADJ_MATS = get_adj_mats(ADJACENCY_DATA, NUM_ENTITIES, NUM_RELATIONS)

            match = np.expand_dims(match, axis=0)

            RULES = [1]
            rel2idx = {1: 1}

            for rule in RULES:
                tableItem = {}
                rule_indices = match[0, :, 1] == rel2idx[rule]
                match_rule = match[:, rule_indices, :]
                preds = model.predict(
                    x=[
                        np.arange(NUM_ENTITIES).reshape(1, -1),
                        match_rule[:, :, 0],
                        match_rule[:, :, 1],
                        match_rule[:, :, 2],
                        ADJ_MATS
                    ]
                )
                y_prob = preds[0]
                tableItem = {
                    'disease_name': disease_name,
                    'drug_name': drug_name,
                    'y_prob': str(y_prob[0]),
                }
                tableData.append(tableItem)
                msg = f'Success! The probability of association between the parasitic disease {disease_name} and the drug {drug_name} is: {y_prob[0]}.'
                logging.info(msg)
                log_messages.append(msg)

        return output.getvalue(), tableData, log_messages

    except Exception as e:
        msg = f"Error in start function: {e}"
        logging.error(msg)
        log_messages.append(msg)
        traceback.print_exc(file=output)
        return output.getvalue(), None, log_messages

if __name__ == '__main__':
    print(start("../uploads/sample.csv"))
